package com.travelmood.controller;
import com.alibaba.fastjson.JSONObject;
import com.travelmood.utils.SignatureUtil;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Mono;
import com.alibaba.fastjson.JSON;


import java.util.Collections;
import java.util.HashMap;
import java.util.Map;

@RestController
@RequestMapping("/image")
public class ImageConvertController {

    private final WebClient webClient;
    private final SignatureUtil signatureUtil;

    @Autowired
    public ImageConvertController(WebClient.Builder webClientBuilder, SignatureUtil signatureUtil) {
        this.webClient = webClientBuilder.baseUrl("http://api-ai.vivo.com.cn").build();
        this.signatureUtil = signatureUtil;
    }

    @PostMapping("/convert")
    public Mono<ResponseEntity<String>> convertImage(@RequestBody Map<String, String> request) {
        String imageBase64 = request.get("imageBase64");
        String prompt = request.get("prompt");

        // 定义请求体数据
        Map<String, Object> data = new HashMap<>();
        data.put("height", 768);
        data.put("width", 576);
        data.put("prompt", prompt);
        data.put("styleConfig", "55c682d5eeca50d4806fd1cba3628781");
        data.put("imageBase64", imageBase64);

        String requestBody = JSON.toJSONString(data);

        return webClient.post()
                .uri("/v1/task_submit")
                .headers(headers -> {
                    headers.setAll(signatureUtil.generateAuthHeaders("POST", "/v1/task_submit", new HashMap<>()));
                    headers.setContentType(MediaType.APPLICATION_JSON);
                })
                .bodyValue(requestBody)
                .retrieve()
                .bodyToMono(String.class)
                .flatMap(body -> {
                    // 解析响应体，提取 task_id
                    String taskId = extractTaskId(body);
                    return checkTaskProgress(taskId);
                })
                .map(url -> ResponseEntity.ok().body(url))
                .onErrorResume(e -> {
                    e.printStackTrace();
                    return Mono.just(ResponseEntity.status(HttpStatus.INTERNAL_SERVER_ERROR).body("Error processing your request: " + e.getMessage()));
                });
    }

    private String extractTaskId(String responseBody) {
        // 解析响应体，提取 task_id
        // 假设响应体中包含一个 "task_id" 字段
        JSONObject jsonObject = JSON.parseObject(responseBody);
        return ((JSONObject) jsonObject).getString("task_id");
    }

    private Mono<String> checkTaskProgress(String taskId) {
        return webClient.get()
                .uri(uriBuilder -> uriBuilder.path("/api/v1/task_progress")
                        .queryParam("task_id", taskId)
                        .build())
                .headers(headers -> {
                    headers.setAll(signatureUtil.generateAuthHeaders("GET", "/api/v1/task_progress", Collections.singletonMap("task_id", taskId)));
                    headers.setContentType(MediaType.APPLICATION_JSON);
                })
                .retrieve()
                .bodyToMono(String.class)
                .flatMap(responseBody -> {
                    // 解析响应体，提取图片 URL
                    String imageUrl = extractImageUrl(responseBody);
                    return Mono.just(imageUrl);
                });
    }

    private String extractImageUrl(String responseBody) {
        // 解析响应体，提取 image_url
        // 假设响应体中包含一个 "image_url" 字段
        JSONObject jsonObject = JSON.parseObject(responseBody);
        return jsonObject.getString("image_url");
    }
}

